#!/usr/bin/env python
# -*- coding=utf-8 -*-
"""
@author: xingwg
@license: (C) Copyright 2020-2025.
@contact: xingweiguo@chinasvt.com
@project: boya-reid
@file: export_graph.py
@time: 2020/10/3 16:00
@desc:
"""
import os
import argparse
import torch
from src.config import cfg
from src.dataset import make_dataloader
from src.model.make_model import make_model

parser = argparse.ArgumentParser(description="ReID Baseline Training")
parser.add_argument(
    "--config_file",
    default="",
    help="path to config file",
    type=str
)
parser.add_argument(
    "opts",
    help="Modify config options using the command-line",
    default=None,
    nargs=argparse.REMAINDER
)

args = parser.parse_args()

if args.config_file != "":
    cfg.merge_from_file(args.config_file)
cfg.merge_from_list(args.opts)
cfg.freeze()

dummy_input = torch.randn(1, 3, cfg.INPUT.SIZE_TEST[0], cfg.INPUT.SIZE_TEST[1], device='cuda')

os.environ["CUDA_VISIBLE_DEVICES"] = "0"
device = "cuda"
train_loader, val_loader, num_query, num_classes = make_dataloader(cfg)
model = make_model(cfg, num_classes)
model.load_param(cfg.TEST.WEIGHT)
model.to(device)

# input_names = ["data"] + ["learned_%d" % i for i in range(693, 1163)]
input_names = ["data"]
output_names = ["feat"]

torch.onnx.export(
    model,
    dummy_input,
    "{}.onnx".format(cfg.MODEL.NAME),
    verbose=True,
    input_names=input_names,
    output_names=output_names,
    # opset_version=7,
    # export_params=True
)
